from .order_book import OrderBook
from .row_data_csv_reader import RowDataCSVReader
from .stock_order import StockOrder
from .snapshot import Snapshot
import os
import csv

class BackTest:
    def __init__(self, order_book: OrderBook, order_csv_path: str, log_path="logs/backtest_price.csv"):
        self.__order_book = order_book
        self.__csv_reader = RowDataCSVReader(order_csv_path)
        self.log_path = log_path
        self.init()
        
    
    def init(self) :
        os.makedirs(os.path.dirname(self.log_path), exist_ok=True)
        # 初始化 CSV 文件，只写入时间戳和价格
        with open(self.log_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["timestamp", "price", "total_trade_volume", "total_trade_value", "buy1_price", "buy1_volume", 
                             "buy2_price", "buy2_volume", "buy3_price", "buy3_volume", "buy4_price", "buy4_volume", "buy5_price", "buy5_volume", 
                             "sell1_price", "sell1_volume", "sell2_price", "sell2_volume", "sell3_price", "sell3_volume", "sell4_price", "sell4_volume", 
                             "sell5_price", "sell5_volume"])
        # 读取竞价数据并进行竞价撮合
        self.load_auction_orderbook()
        self._save_price()
    

    def load_auction_orderbook(self):
        #读取9点半之前的数据，只添加到订单簿中，不进行撮合
        auction_rows, ts = self.__csv_reader.read_before_930()
        for row_data in auction_rows:
            order = row_data.to_stock_order()
            if row_data.is_cancel():
                self.__order_book.cancel_order(order)
            else:
                self.__order_book.add_to_book(order)
        # 竞价
        self.__order_book.auction()
        self.__order_book.set_timestamp(ts)

        #打印当前订单簿快照
        # snapshot: Snapshot = self.__order_book.get_snapshot()  # 获取当前快照
        # print("Auction Snapshot:")
        # print(snapshot)
        # print("-------------------------------------")

    def get_order_book(self):
        return self.__order_book

    def run(self):
        """
        执行回测
        """
        if self.__csv_reader is None:
            raise ValueError("CSV reader is not set. Please call set_csv_reader() first.")

        while True:
            # 读取下一秒的数据
            row_datas, end_time, is_over = self.__csv_reader.read_next_n_seconds(self.__order_book.get_timestamp(), 3)
            if is_over:
                break
            for row_data in row_datas:
                # 将行数据转换为订单并添加到订单簿
                order = row_data.to_stock_order()
                # 判断是否为撤单
                if row_data.is_cancel():
                    self.__order_book.cancel_order(order)
                else:   
                    self.__order_book.add_order_and_match(order)
            self.__order_book.set_timestamp(end_time)
             # 保存最新价格和总成交量、总成交额
            self._save_price()
            snapshot: Snapshot = self.__order_book.get_snapshot()  # 获取当前快照
            print(snapshot)
            print("-------------------------------------")

    def step_once(self):
        """
        执行一次回测步骤，读取下一行数据并更新订单簿
        """
        row_datas, end_time, is_over = self.__csv_reader.read_next_n_seconds(self.__order_book.get_timestamp(), 3)
        if is_over:
            return
        for row_data in row_datas:
            # 将行数据转换为订单并添加到订单簿
            order = row_data.to_stock_order()
            # 判断是否为撤单
            if row_data.is_cancel():
                self.__order_book.cancel_order(order)
            else:   
                self.__order_book.add_order_and_match(order)
        self.__order_book.set_timestamp(end_time)
        self._save_price()
        # snapshot: Snapshot = self.__order_book.get_snapshot()  # 获取当前快照
        # print(snapshot)
        return
    
    def _save_price(self):
        """保存最新价格和5档行情"""
        with open(self.log_path, "a", newline="") as f:
            writer = csv.writer(f)

            # 获取基础信息
            timestamp = self.__order_book.get_timestamp()
            price = self.__order_book.get_price()
            total_volume = self.__order_book.get_total_trade_volume()
            total_value = self.__order_book.get_total_trade_value()

            # 获取订单簿快照（包含10档数据，我们取前5档）
            snapshot = self.__order_book.get_snapshot()

            # 提取买1-买5档的价格和数量
            buy_prices = snapshot.get_buy_prices()[:5]  # 前5档买价
            buy_volumes = snapshot.get_buy_volumes()[:5]  # 前5档买量

            # 提取卖1-卖5档的价格和数量
            sell_prices = snapshot.get_sell_prices()[:5]  # 前5档卖价
            sell_volumes = snapshot.get_sell_volumes()[:5]  # 前5档卖量

            # 构建行数据：基础信息 + 5档买卖盘
            row_data = [
                timestamp, 
                price, 
                total_volume, 
                total_value,
            ]

            # 添加买1-买5档数据
            for i in range(5):
                row_data.extend([buy_prices[i], buy_volumes[i]])

            # 添加卖1-卖5档数据
            for i in range(5):
                row_data.extend([sell_prices[i], sell_volumes[i]])

            writer.writerow(row_data)
    
    def place_order(self, order:StockOrder):
        self.__order_book.add_order_and_match(order)
        self._save_price()
        return 
    
    def is_done(self) -> bool:
        """
        检查回测是否完成
        :return: 如果没有更多数据可读,则返回True
        """
        return self.__csv_reader.get_current_position() >= self.__csv_reader.get_total_rows()
    
    def reset(self):
        """
        重置回测状态
        """
        self.__csv_reader.reset()
        self.__order_book.clear()  
        self.init()


            